{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "功能：训练学生模型。\n",
    "将教师模型设置为评估模式以确保在训练学生模型时不进行参数更新。\n",
    "将学生模型设置为训练模式以启用梯度计算和参数更新。\n",
    "使用训练数据加载器来按批次加载训练数据和对应的标签。\n",
    "对于每一个批次的数据和标签，执行以下操作：\n",
    "前向传播，获取教师模型和学生模型的输出\n",
    "计算损失，包含分类损失和知识蒸馏损失\n",
    "反向传播\n",
    "参数更新"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "整体目的  \n",
    "│  \n",
    "├── 训练学生模型 (使用知识蒸馏)  \n",
    "│  \n",
    "┌───────────────────────────────┐  \n",
    "│         输入部分             │  \n",
    "│ ┌───────────────┐           │  \n",
    "│ │ teacher_model │           │  \n",
    "│ ├───────────────┤           │  \n",
    "│ │ student_model │           │  \n",
    "│ ├───────────────┤           │  \n",
    "│ │   optimizer   │           │  \n",
    "│ ├───────────────┤           │  \n",
    "│ │   criterion   │           │  \n",
    "│ ├───────────────┤           │  \n",
    "│ │ distillation_loss │       │  \n",
    "│ ├───────────────┤           │  \n",
    "│ │   train_loader │           │  \n",
    "│ └───────────────┘           │  \n",
    "└───────────────────────────────┘  \n",
    "│  \n",
    "┌───────────────────────────────┐  \n",
    "│         训练过程             │  \n",
    "│ ┌───────────────┐           │  \n",
    "│ │ 设置模型模式  │           │  \n",
    "│ │ teacher_model.eval() │    │  \n",
    "│ │ student_model.train() │    │  \n",
    "│ └───────────────┘           │  \n",
    "│      │      │                 │  \n",
    "│      └──────┼─────────────────┐  \n",
    "│             └───foreach batch──┐  \n",
    "│                          ┌─────────────┐  \n",
    "│                          │ 梯度清零   │  \n",
    "│                          │ optimizer.zero_grad() │  \n",
    "│                          └─────────────┘  \n",
    "│                              ├─────────────┐  \n",
    "│ ┌────────────────────────────┴─────────────┐  \n",
    "│ │            教师模型输出                   │  \n",
    "│ │ teacher_output = teacher_model(images.view(images.size(0),-1)) │  \n",
    "│ └────────────────────────────┬─────────────┘  \n",
    "│                              │  \n",
    "│ ┌─────────────────────────────┐  \n",
    "│ │            学生模型输出    │  \n",
    "│ │ student_output = student_model(images.view(images.size(0),-1)) │  \n",
    "│ └─────────────────────────────┘  \n",
    "│         ┌─────────────┐  \n",
    "│         │ 损失计算    │  \n",
    "│ ┌─────────────────────────────┐  \n",
    "│ │ loss = criterion(student_output, labels) │  \n",
    "│ ├─────────────────────────────┤  \n",
    "│ │ loss_distillation = distillation_loss(    │  \n",
    "│ │     nn.functional.log_softmax(student_output / 2, dim=1),  │  \n",
    "│ │     nn.functional.softmax(teacher_output / 2, dim=1) │  \n",
    "│ │ )                                     │  \n",
    "│ ├─────────────────────────────┤  \n",
    "│ │ total_loss = loss + loss_distillation │  \n",
    "│ └─────────────────────────────┘  \n",
    "│         │  \n",
    "│         └─────────────┐  \n",
    "│                       │  \n",
    "│ ┌─────────────────────────────┐  \n",
    "│ │  计算梯度和更新参数        │  \n",
    "│ │ total_loss.backward()       │  \n",
    "│ │ optimizer.step()            │  \n",
    "│ └─────────────────────────────┘  \n",
    "└───────────────────────────────┘  \n",
    "│  \n",
    "┌───────────────────────────────┐  \n",
    "│         模型保存             │  \n",
    "│ torch.save(student_model.state_dict(), model_path) │  \n",
    "└───────────────────────────────┘"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np   \n",
    "import torch  \n",
    "import torch.nn as nn  \n",
    "\n",
    "def train_student(teacher_model, student_model, optimizer, criterion, distillation_loss, train_loader, model_path=\"student_model.pth\"):  \n",
    "    # 设置教师模型为评估模式，以便在推断时禁用某些特性（如 dropout 和 batch normalization）。  \n",
    "    teacher_model.eval()  \n",
    "    \n",
    "    # 设置学生模型为训练模式，使其能够进行参数更新。  \n",
    "    student_model.train()  \n",
    "    \n",
    "    # 遍历训练数据加载器，获取每个批次的图像和标签  \n",
    "    for images, labels in train_loader:  \n",
    "        # 清空梯度以便进行新一轮的训练  \n",
    "        optimizer.zero_grad()  \n",
    "\n",
    "        # 使用上下文管理器来禁用梯度计算，以提高推断的效率  \n",
    "        with torch.no_grad():  \n",
    "            # 获取教师模型对当前批次图像的预测输出  \n",
    "            teacher_output = teacher_model(images.view(images.size(0), -1))  \n",
    "\n",
    "        # 通过学生模型获取当前批次图像的预测输出  \n",
    "        student_output = student_model(images.view(images.size(0), -1))  \n",
    "\n",
    "        # 计算学生模型输出与真实标签之间的损失  \n",
    "        loss = criterion(student_output, labels)  \n",
    "\n",
    "        # 计算知识蒸馏损失，这里使用了温度缩放  \n",
    "        loss_distillation = distillation_loss(  \n",
    "            nn.functional.log_softmax(student_output / 2, dim=1),  # 学生模型的输出经过 softmax  \n",
    "            nn.functional.softmax(teacher_output / 2, dim=1)        # 教师模型的输出经过 softmax  \n",
    "        )  \n",
    "\n",
    "        # 总损失是真实标签损失和知识蒸馏损失之和  \n",
    "        total_loss = loss + loss_distillation  \n",
    "\n",
    "        # 反向传播，计算每个参数的梯度  \n",
    "        total_loss.backward()  \n",
    "\n",
    "        # 更新学生模型的参数  \n",
    "        optimizer.step()  \n",
    "\n",
    "    # 训练完成后，将学生模型的状态字典保存到指定的路径  \n",
    "    torch.save(student_model.state_dict(), model_path)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorchgpu",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
